{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import sys, subprocess, json, os\n",
        "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"dagster\", \"pandas\", \"scikit-learn\"])\n",
        "\n",
        "import numpy as np, pandas as pd\n",
        "from pathlib import Path\n",
        "from dagster import (\n",
        "    asset, AssetCheckResult, asset_check, Definitions, materialize, Output,\n",
        "    DailyPartitionsDefinition, IOManager, io_manager\n",
        ")\n",
        "from sklearn.linear_model import LinearRegression\n",
        "\n",
        "BASE = Path(\"/content/dagstore\"); BASE.mkdir(parents=True, exist_ok=True)\n",
        "START = \"2025-08-01\""
      ],
      "metadata": {
        "id": "M5Qsa0qWBlCi"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CSVIOManager(IOManager):\n",
        "    def __init__(self, base: Path): self.base = base\n",
        "    def _path(self, key, ext): return self.base / f\"{'_'.join(key.path)}.{ext}\"\n",
        "    def handle_output(self, context, obj):\n",
        "        if isinstance(obj, pd.DataFrame):\n",
        "            p = self._path(context.asset_key, \"csv\"); obj.to_csv(p, index=False)\n",
        "            context.log.info(f\"Saved {context.asset_key} -> {p}\")\n",
        "        else:\n",
        "            p = self._path(context.asset_key, \"json\"); p.write_text(json.dumps(obj, indent=2))\n",
        "            context.log.info(f\"Saved {context.asset_key} -> {p}\")\n",
        "    def load_input(self, context):\n",
        "        k = context.upstream_output.asset_key; p = self._path(k, \"csv\")\n",
        "        df = pd.read_csv(p); context.log.info(f\"Loaded {k} <- {p} ({len(df)} rows)\"); return df\n",
        "\n",
        "@io_manager\n",
        "def csv_io_manager(_): return CSVIOManager(BASE)\n",
        "\n",
        "daily = DailyPartitionsDefinition(start_date=START)"
      ],
      "metadata": {
        "id": "o-R2uk2qBoov"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@asset(partitions_def=daily, description=\"Synthetic raw sales with noise & occasional nulls.\")\n",
        "def raw_sales(context) -> Output[pd.DataFrame]:\n",
        "    rng = np.random.default_rng(42)\n",
        "    n = 200; day = context.partition_key\n",
        "    x = rng.normal(100, 20, n); promo = rng.integers(0, 2, n); noise = rng.normal(0, 10, n)\n",
        "    sales = 2.5 * x + 30 * promo + noise + 50\n",
        "    x[rng.choice(n, size=max(1, n // 50), replace=False)] = np.nan\n",
        "    df = pd.DataFrame({\"date\": day, \"units\": x, \"promo\": promo, \"sales\": sales})\n",
        "    meta = {\"rows\": n, \"null_units\": int(df[\"units\"].isna().sum()), \"head\": df.head().to_markdown()}\n",
        "    return Output(df, metadata=meta)\n",
        "\n",
        "@asset(description=\"Clean nulls, clip outliers for robust downstream modeling.\")\n",
        "def clean_sales(context, raw_sales: pd.DataFrame) -> Output[pd.DataFrame]:\n",
        "    df = raw_sales.dropna(subset=[\"units\"]).copy()\n",
        "    lo, hi = df[\"units\"].quantile([0.01, 0.99]); df[\"units\"] = df[\"units\"].clip(lo, hi)\n",
        "    meta = {\"rows\": len(df), \"units_min\": float(df.units.min()), \"units_max\": float(df.units.max())}\n",
        "    return Output(df, metadata=meta)\n",
        "\n",
        "@asset(description=\"Feature engineering: interactions & standardized columns.\")\n",
        "def features(context, clean_sales: pd.DataFrame) -> Output[pd.DataFrame]:\n",
        "    df = clean_sales.copy()\n",
        "    df[\"units_sq\"] = df[\"units\"] ** 2; df[\"units_promo\"] = df[\"units\"] * df[\"promo\"]\n",
        "    for c in [\"units\", \"units_sq\", \"units_promo\"]:\n",
        "        mu, sigma = df[c].mean(), df[c].std(ddof=0) or 1.0\n",
        "        df[f\"z_{c}\"] = (df[c] - mu) / sigma\n",
        "    return Output(df, metadata={\"rows\": len(df), \"cols\": list(df.columns)})"
      ],
      "metadata": {
        "id": "yP1KsT6GBvfK"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@asset_check(asset=clean_sales, description=\"No nulls; promo in {0,1}; units within clipped bounds.\")\n",
        "def clean_sales_quality(clean_sales: pd.DataFrame) -> AssetCheckResult:\n",
        "    nulls = int(clean_sales.isna().sum().sum())\n",
        "    promo_ok = bool(set(clean_sales[\"promo\"].unique()).issubset({0, 1}))\n",
        "    units_ok = bool(clean_sales[\"units\"].between(clean_sales[\"units\"].min(), clean_sales[\"units\"].max()).all())\n",
        "    passed = bool((nulls == 0) and promo_ok and units_ok)\n",
        "    return AssetCheckResult(\n",
        "        passed=passed,\n",
        "        metadata={\"nulls\": nulls, \"promo_ok\": promo_ok, \"units_ok\": units_ok},\n",
        "    )\n",
        "\n",
        "@asset(description=\"Train a tiny linear regressor; emit R^2 and coefficients.\")\n",
        "def tiny_model_metrics(context, features: pd.DataFrame) -> dict:\n",
        "    X = features[[\"z_units\", \"z_units_sq\", \"z_units_promo\", \"promo\"]].values\n",
        "    y = features[\"sales\"].values\n",
        "    model = LinearRegression().fit(X, y)\n",
        "    return {\"r2_train\": float(model.score(X, y)),\n",
        "            **{n: float(c) for n, c in zip([\"z_units\",\"z_units_sq\",\"z_units_promo\",\"promo\"], model.coef_)}}"
      ],
      "metadata": {
        "id": "2SaNVSS8By0z"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zNEIXqo36FH7",
        "outputId": "af9df168-9feb-4074-fe61-d174e832468d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - RUN_START - Started execution of run for \"__ephemeral_asset_job__\".\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - ENGINE_EVENT - Executing steps in process (pid: 479)\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - RESOURCE_INIT_STARTED - Starting initialization of resources [io_manager].\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - RESOURCE_INIT_SUCCESS - Finished initialization of resources [io_manager].\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - LOGS_CAPTURED - Started capturing logs in process (pid: 479).\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - raw_sales - STEP_START - Started execution of step \"raw_sales\".\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - raw_sales - STEP_OUTPUT - Yielded output \"result\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - raw_sales - Saved AssetKey(['raw_sales']) -> /content/dagstore/raw_sales.csv\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - raw_sales - ASSET_MATERIALIZATION - Materialized value raw_sales.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - raw_sales - HANDLED_OUTPUT - Handled output \"result\" using IO manager \"io_manager\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - raw_sales - STEP_SUCCESS - Finished execution of step \"raw_sales\" in 29ms.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - STEP_START - Started execution of step \"clean_sales\".\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Materializing everything for: 2025-08-01\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - clean_sales - Loaded AssetKey(['raw_sales']) <- /content/dagstore/raw_sales.csv (200 rows)\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - LOADED_INPUT - Loaded input \"raw_sales\" using input manager \"io_manager\", from output \"result\" of step \"raw_sales\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - STEP_INPUT - Got input \"raw_sales\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - STEP_OUTPUT - Yielded output \"result\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - clean_sales - Saved AssetKey(['clean_sales']) -> /content/dagstore/clean_sales.csv\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - ASSET_MATERIALIZATION - Materialized value clean_sales.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - HANDLED_OUTPUT - Handled output \"result\" using IO manager \"io_manager\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales - STEP_SUCCESS - Finished execution of step \"clean_sales\" in 58ms.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - STEP_START - Started execution of step \"clean_sales_clean_sales_quality\".\n",
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - clean_sales_clean_sales_quality - Loaded AssetKey(['clean_sales']) <- /content/dagstore/clean_sales.csv (196 rows)\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - LOADED_INPUT - Loaded input \"clean_sales\" using input manager \"io_manager\", from output \"result\" of step \"clean_sales\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - STEP_INPUT - Got input \"clean_sales\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - STEP_OUTPUT - Yielded output \"result\" of type \"Any\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - ASSET_CHECK_EVALUATION - Asset check 'clean_sales_quality' on 'clean_sales' passed.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - clean_sales_clean_sales_quality - STEP_SUCCESS - Finished execution of step \"clean_sales_clean_sales_quality\" in 34ms.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - STEP_START - Started execution of step \"features\".\n",
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - features - Loaded AssetKey(['clean_sales']) <- /content/dagstore/clean_sales.csv (196 rows)\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - LOADED_INPUT - Loaded input \"clean_sales\" using input manager \"io_manager\", from output \"result\" of step \"clean_sales\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - STEP_INPUT - Got input \"clean_sales\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - STEP_OUTPUT - Yielded output \"result\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:00 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - features - Saved AssetKey(['features']) -> /content/dagstore/features.csv\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - ASSET_MATERIALIZATION - Materialized value features.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - HANDLED_OUTPUT - Handled output \"result\" using IO manager \"io_manager\"\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - features - STEP_SUCCESS - Finished execution of step \"features\" in 43ms.\n",
            "2025-08-16 04:03:00 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - STEP_START - Started execution of step \"tiny_model_metrics\".\n",
            "2025-08-16 04:03:01 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - tiny_model_metrics - Loaded AssetKey(['features']) <- /content/dagstore/features.csv (196 rows)\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - LOADED_INPUT - Loaded input \"features\" using input manager \"io_manager\", from output \"result\" of step \"features\"\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - STEP_INPUT - Got input \"features\" of type \"DataFrame\". (Type check passed).\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - STEP_OUTPUT - Yielded output \"result\" of type \"dict\". (Type check passed).\n",
            "2025-08-16 04:03:01 +0000 - dagster - INFO - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - tiny_model_metrics - Saved AssetKey(['tiny_model_metrics']) -> /content/dagstore/tiny_model_metrics.json\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - ASSET_MATERIALIZATION - Materialized value tiny_model_metrics.\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - HANDLED_OUTPUT - Handled output \"result\" using IO manager \"io_manager\"\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - tiny_model_metrics - STEP_SUCCESS - Finished execution of step \"tiny_model_metrics\" in 43ms.\n",
            "/usr/local/lib/python3.11/dist-packages/dagster/_core/execution/context_creation_job.py:265: RuntimeWarning: coroutine 'BaseEventLoop.shutdown_asyncgens' was never awaited\n",
            "  pass\n",
            "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - ENGINE_EVENT - Finished steps in process (pid: 479) in 262ms\n",
            "2025-08-16 04:03:01 +0000 - dagster - DEBUG - __ephemeral_asset_job__ - 52884413-5e11-4434-8a98-3e44dc808c28 - 479 - RUN_SUCCESS - Finished execution of run for \"__ephemeral_asset_job__\".\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Run success: True\n",
            "raw_sales.csv -> 9894 bytes\n",
            "clean_sales.csv -> 9731 bytes\n",
            "features.csv -> 27106 bytes\n",
            "tiny_model_metrics.json -> 174 bytes\n",
            "Metrics: {'r2_train': 0.9547421316644358, 'z_units': 32.931720119978166, 'z_units_sq': 11.998171264446981, 'z_units_promo': -6.324476924103006, 'promo': 44.90030628506354}\n"
          ]
        }
      ],
      "source": [
        "defs = Definitions(\n",
        "    assets=[raw_sales, clean_sales, features, tiny_model_metrics, clean_sales_quality],\n",
        "    resources={\"io_manager\": csv_io_manager}\n",
        ")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    run_day = os.environ.get(\"RUN_DATE\") or START\n",
        "    print(\"Materializing everything for:\", run_day)\n",
        "    result = materialize(\n",
        "        [raw_sales, clean_sales, features, tiny_model_metrics, clean_sales_quality],\n",
        "        partition_key=run_day,\n",
        "        resources={\"io_manager\": csv_io_manager},\n",
        "    )\n",
        "    print(\"Run success:\", result.success)\n",
        "\n",
        "    for fname in [\"raw_sales.csv\",\"clean_sales.csv\",\"features.csv\",\"tiny_model_metrics.json\"]:\n",
        "        f = BASE / fname\n",
        "        if f.exists():\n",
        "            print(fname, \"->\", f.stat().st_size, \"bytes\")\n",
        "            if fname.endswith(\".json\"):\n",
        "                print(\"Metrics:\", json.loads(f.read_text()))"
      ]
    }
  ]
}